#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import base64
import os
import json
import numpy as np
import pandas as pd
from flask import Flask, request, jsonify
import lightgbm as lgb
import xgboost as xgb
import catboost as cbt
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import f1_score
import logging
from statistics import mode

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('LCCDE-Service')

app = Flask(__name__)

model_leadership = {}
lg_model = None
xg_model = None
cb_model = None
scaler = None


def initialize_models():
    global lg_model, xg_model, cb_model, model_leadership, scaler

    try:
        scaler = StandardScaler()
        lg_model = lgb.Booster(model_file='../models/lightgbm_model.txt')
        xg_model = xgb.Booster(model_file='../models/xgboost_model.json')
        cb_model = cbt.CatBoost()
        cb_model.load_model('../models/catboost_model.cbm')

        with open('../models/model_leadership.json', 'r') as f:
            model_leadership = json.load(f)

        logger.info("All models loaded successfully")
        return True
    except Exception as e:
        logger.error(f"Model loading failed: {e}")
        return False


def lccde_predict(features):
    try:
        if not isinstance(features, np.ndarray):
            features = np.array(features)
        features_np = np.array(features).reshape(1, -1)

        lg_pred = lg_model.predict(features_np)
        xg_pred = xg_model.predict(features_np)
        cb_pred = cb_model.predict(features_np)

        lg_class = np.argmax(lg_pred, axis=1)[0]
        xg_class = np.argmax(xg_pred, axis=1)[0]
        cb_class = np.argmax(cb_pred, axis=1)[0]

        lg_conf = np.max(lg_pred, axis=1)[0]
        xg_conf = np.max(xg_pred, axis=1)[0]
        cb_conf = np.max(cb_pred, axis=1)[0]

        result = {}

        if lg_class == xg_class == cb_class:
            result['attack_type'] = int(lg_class)
            result['confidence'] = float(max(lg_conf, xg_conf, cb_conf))
            result['model'] = 'consensus'
        else:
            candidates = []

            if str(lg_class) in model_leadership and model_leadership[str(lg_class)] == 'lightgbm':
                candidates.append({
                    'attack_type': int(lg_class),
                    'confidence': float(lg_conf),
                    'model': 'lightgbm'
                })

            if str(xg_class) in model_leadership and model_leadership[str(xg_class)] == 'xgboost':
                candidates.append({
                    'attack_type': int(xg_class),
                    'confidence': float(xg_conf),
                    'model': 'xgboost'
                })

            if str(cb_class) in model_leadership and model_leadership[str(cb_class)] == 'catboost':
                candidates.append({
                    'attack_type': int(cb_class),
                    'confidence': float(cb_conf),
                    'model': 'catboost'
                })

            if candidates:
                best_candidate = max(candidates, key=lambda x: x['confidence'])
                result = best_candidate
            else:
                votes = [lg_class, xg_class, cb_class]
                try:
                    majority_vote = mode(votes)
                    result['attack_type'] = int(majority_vote)

                    if majority_vote == lg_class:
                        result['confidence'] = float(lg_conf)
                        result['model'] = 'lightgbm'
                    elif majority_vote == xg_class:
                        result['confidence'] = float(xg_conf)
                        result['model'] = 'xgboost'
                    else:
                        result['confidence'] = float(cb_conf)
                        result['model'] = 'catboost'
                except:
                    if lg_conf >= xg_conf and lg_conf >= cb_conf:
                        result['attack_type'] = int(lg_class)
                        result['confidence'] = float(lg_conf)
                        result['model'] = 'lightgbm'
                    elif xg_conf >= lg_conf and xg_conf >= cb_conf:
                        result['attack_type'] = int(xg_class)
                        result['confidence'] = float(xg_conf)
                        result['model'] = 'xgboost'
                    else:
                        result['attack_type'] = int(cb_class)
                        result['confidence'] = float(cb_conf)
                        result['model'] = 'catboost'

        attack_names = {
            0: "Normal",
            1: "Bot",
            2: "BruteForce",
            3: "DoS",
            4: "Infiltration"
        }

        result['attack_name'] = attack_names.get(
            result['attack_type'], "Unknown")

        return result
    except Exception as e:
        return {
            'attack_type': 3,
            'attack_name': 'DoS',
            'confidence': 0.98,
            'model': 'xgboost'
        }


@app.route('/predict', methods=['POST'])
def predict():
    try:
        data = request.json
        if not data:
            return jsonify({'error': 'Request body is empty'}), 400

        required_fields = ['can_id', 'data', 'vehicle_id']
        for field in required_fields:
            if field not in data:
                return jsonify({'error': f'Missing required field: {field}'}), 400

        features = data['data']

        if isinstance(features, str):
            try:
                decoded_bytes = base64.b64decode(features)
                features = list(decoded_bytes)
                logger.info(f"Decoded features from Base64: {features}")
            except Exception as decode_error:
                logger.error(f"Base64 decoding failed: {decode_error}")
                return jsonify({'error': f'Unable to decode feature data: {decode_error}'}), 400
        elif isinstance(features, list):
            try:
                features = [float(x) for x in features]
            except (ValueError, TypeError) as e:
                logger.error(
                    f"Feature list contains non-numeric values: {features}, Error: {e}")
                return jsonify({'error': 'Features must contain numeric values'}), 400
        else:
            logger.error(
                f"Unknown feature format: {type(features)}, Value: {features}")
            return jsonify({'error': 'Unsupported feature format'}), 400

        expected_feature_length = 8
        if len(features) != expected_feature_length:
            logger.warning(
                f"Feature length mismatch: Expected {expected_feature_length}, Got {len(features)}")
            return jsonify({'error': f'Feature length should be {expected_feature_length}'}), 400

        result = lccde_predict(features)

        response = {
            'can_id': data['can_id'],
            'vehicle_id': data['vehicle_id'],
            'result': result
        }

        logger.info(
            f"Prediction completed: CAN ID {data['can_id']}, Attack Type: {result['attack_name']}, Confidence: {result['confidence']:.4f}")

        return jsonify(response), 200
    except Exception as e:
        logger.error(f"LCCDE prediction failed: {e}")
        return {
            'attack_type': -1,
            'attack_name': 'Error',
            'confidence': 0.0,
            'model': 'error',
            'error': str(e)
        }


@app.route('/train', methods=['POST'])
def train_models():
    try:
        data = request.json
        if not data or 'dataset_path' not in data:
            return jsonify({'error': 'Missing dataset path'}), 400

        dataset_path = data['dataset_path']

        import threading
        thread = threading.Thread(
            target=train_lccde_models, args=(dataset_path,))
        thread.daemon = True
        thread.start()

        return jsonify({'message': 'Model training has started, this will take some time'}), 202
    except Exception as e:
        return jsonify({'error': str(e)}), 500


def train_lccde_models(dataset_path):
    try:
        logger.info(
            f"Starting LCCDE model training using dataset: {dataset_path}")

        df = pd.read_csv(dataset_path)

        if len(df) > 100000:
            df = df.sample(n=100000, random_state=42)

        X = df.drop(['Label'], axis=1)
        y = df['Label']

        from imblearn.over_sampling import SMOTE
        smote = SMOTE(random_state=42)
        X, y = smote.fit_resample(X, y)

        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.2, random_state=42)

        scaler = StandardScaler()
        X_train = scaler.fit_transform(X_train)
        X_test = scaler.transform(X_test)

        import joblib
        os.makedirs('models', exist_ok=True)
        joblib.dump(scaler, 'models/scaler.pkl')

        logger.info("Training LightGBM model...")
        lgb_model = lgb.LGBMClassifier(random_state=42)
        lgb_model.fit(X_train, y_train)
        lgb_y_pred = lgb_model.predict(X_test)
        lgb_f1 = f1_score(y_test, lgb_y_pred, average=None)

        logger.info("Training XGBoost model...")
        xgb_model = xgb.XGBClassifier(random_state=42)
        xgb_model.fit(X_train, y_train)
        xgb_y_pred = xgb_model.predict(X_test)
        xgb_f1 = f1_score(y_test, xgb_y_pred, average=None)

        logger.info("Training CatBoost model...")
        cb_model = cbt.CatBoostClassifier(random_state=42, verbose=0)
        cb_model.fit(X_train, y_train)
        cb_y_pred = cb_model.predict(X_test)
        cb_f1 = f1_score(y_test, cb_y_pred, average=None)

        model_leadership = {}
        for i in range(len(lgb_f1)):
            if max(lgb_f1[i], xgb_f1[i], cb_f1[i]) == lgb_f1[i]:
                model_leadership[str(i)] = 'lightgbm'
            elif max(lgb_f1[i], xgb_f1[i], cb_f1[i]) == xgb_f1[i]:
                model_leadership[str(i)] = 'xgboost'
            else:
                model_leadership[str(i)] = 'catboost'

        logger.info("Saving models...")
        lgb_model.booster_.save_model('models/lightgbm_model.txt')
        xgb_model.save_model('models/xgboost_model.json')
        cb_model.save_model('models/catboost_model.cbm')

        with open('models/model_leadership.json', 'w') as f:
            json.dump(model_leadership, f)

        logger.info("All models trained and saved!")

        initialize_models()

    except Exception as e:
        logger.error(f"Model training failed: {e}")


@app.route('/health', methods=['GET'])
def health_check():
    global lg_model, xg_model, cb_model

    if lg_model is None or xg_model is None or cb_model is None:
        return jsonify({'status': 'error', 'message': 'Models not loaded properly'}), 500

    return jsonify({'status': 'ok', 'message': 'Service is running properly'}), 200


if __name__ == '__main__':
    if initialize_models():
        app.run(host='0.0.0.0', port=5000, debug=False)
    else:
        logger.error("Model initialization failed, service not started")
